from util_pose import get_pose_predictor,get_pose_res,get_pose_pil_by_keypoints
from util_parse import get_parse_predictor,get_parse_res,get_parse_zone_by_clotype
from util_flux import process_img_1024
import os

'''
输入 img_path / img_pil 
        check width&height 必须是 384&512 的倍数
    clo_type:str
    clo_path / clo_pil
处理
    img -> pose -> pose_pil --resize to img.size--> pose_pil
    img -> parse -> mask    --resize to img.size--> mask
    masked_img = img.paste( mask, white color)
    condiction_pose = masked_img.paste( mask, pose_pil )

    mask_rect,rect_w,rect_y = 外接矩形(mask) # mask中不规则的区域变成外接矩形状
    clo_resized = clo_pil.resize( rect_w, rect_y )
    condiction_clo = masked_img.paste( mask_rect, clo_resized )

    final_img = concat(img, condiction_clo, condiction_pose)
'''
from PIL import Image, ImageDraw
import numpy as np

def get_masked_img(img_pil, mask_pil):
    """
    Returns a white background image with the masked region from img_pil.
    """
    img = img_pil.convert("RGB")
    mask = mask_pil.convert("L").resize(img.size, Image.NEAREST)
    white_bg = Image.new("RGB", img.size, (255, 255, 255))
    # 反转mask再paste
    inv_mask = mask.point(lambda p: 255 - p)
    white_bg.paste(img, mask=inv_mask)
    return white_bg

def get_condition_pose(masked_img, pose_pil, mask_pil):
    """
    Returns masked_img with pose_pil pasted in the masked region.
    """
    pose_img = pose_pil.convert("RGB").resize(masked_img.size, Image.BILINEAR)
    mask = mask_pil.convert("L").resize(masked_img.size, Image.NEAREST)
    cond_pose = masked_img.copy()
    cond_pose.paste(pose_img, mask=mask)
    return cond_pose

def get_mask_rect(mask_pil):
    """
    Returns the bounding box (left, upper, right, lower) of the nonzero region in mask_pil,
    and the width and height of the rectangle.
    """
    mask = mask_pil.convert("L")
    arr = np.array(mask)
    ys, xs = np.where(arr > 0)
    if len(xs) == 0 or len(ys) == 0:
        # No mask region
        return (0, 0, mask.width, mask.height), mask.width, mask.height
    left, right = xs.min(), xs.max()
    upper, lower = ys.min(), ys.max()
    rect = (left, upper, right+1, lower+1)
    rect_w, rect_h = right - left + 1, lower - upper + 1
    return rect, rect_w, rect_h

def get_condition_clo(masked_img, mask_pil, clo_pil):
    """
    Returns masked_img with the clothing image resized to the mask's bounding box and pasted in that region.
    """
    # 将mask_pil转换为灰度图并调整为masked_img的大小
    mask = mask_pil.convert("L").resize(masked_img.size, Image.NEAREST)
    # 获取mask的外接矩形及其宽高
    rect, rect_w, rect_h = get_mask_rect(mask)
    # 如果没有有效的mask区域，则直接返回原图
    if rect_w <= 0 or rect_h <= 0:
        return masked_img.copy()
    # 将clo_pil缩放到外接矩形的大小
    clo_resized = process_img_1024('',img_pil=clo_pil,target_shape=(rect_w,rect_h))
    # 拷贝一份底图用于粘贴衣服
    cond_clo = masked_img.copy()
    # 从mask中裁剪出外接矩形区域
    mask_rect = mask.crop(rect)
    # 生成二值化的粘贴mask（非零为255，零为0）
    paste_mask = mask_rect.point(lambda p: 255)
    # 创建一个与外接矩形同样大小的白底区域
    # region = Image.new("RGB", (rect_w, rect_h), (255,255,255))
    # 将缩放后的衣服粘贴到region上，仅在paste_mask为255的位置
    # region.paste(clo_resized, mask=paste_mask)
    # 将region粘贴回cond_clo的对应位置，仍然用paste_mask控制透明度
    cond_clo.paste(clo_resized, box=(rect[0], rect[1]), mask=paste_mask)
    # 返回最终合成的图片
    return cond_clo

def concat_images(img_list, direction='horizontal'):
    """
    Concatenate a list of PIL images horizontally or vertically.
    """
    if not img_list:
        return None
    widths, heights = zip(*(img.size for img in img_list))
    if direction == 'horizontal':
        total_width = sum(widths)
        max_height = max(heights)
        new_img = Image.new('RGB', (total_width, max_height), (255,255,255))
        x_offset = 0
        for im in img_list:
            new_img.paste(im, (x_offset, 0))
            x_offset += im.width
    else:
        max_width = max(widths)
        total_height = sum(heights)
        new_img = Image.new('RGB', (max_width, total_height), (255,255,255))
        y_offset = 0
        for im in img_list:
            new_img.paste(im, (0, y_offset))
            y_offset += im.height
    return new_img

def tryon_pipeline(
    img_path=None, img_pil=None,
    clo_type='upper',
    clo_path=None, clo_pil=None,
    clo_mask_path=None, clo_mask_pil=None,
    pose_predictor=None, parse_predictor=None,
    target_shape=(384,512)
):
    """
    Main pipeline for mask try-on.
    """
    # Load image
    if img_pil is None:
        img_pil = Image.open(img_path).convert("RGB")
    img_w, img_h = img_pil.size
    if img_w % target_shape[0] != 0 or img_h % target_shape[1] != 0:
        raise ValueError(f"Image size {img_pil.size} must be a multiple of {target_shape}")

    # Load clothing
    if clo_pil is None:
        clo_pil = Image.open(clo_path).convert("RGB")
    # 额外做一步，读取clo mask, 获得mask中前景的外接矩形, 提取clo_pil的前景
    if clo_mask_pil is None:
        clo_mask_pil = Image.open(clo_mask_path).convert("L")
    clo_mask_rect, clo_mask_w, clo_mask_h = get_mask_rect(clo_mask_pil)
    clo_pil = clo_pil.crop(clo_mask_rect)

    # Get predictors
    if pose_predictor is None:
        pose_predictor = get_pose_predictor()
    if parse_predictor is None:
        parse_predictor = get_parse_predictor()

    # Get pose and pose PIL
    pose_keypoints = get_pose_res(pose_predictor, img_pil=img_pil, pose_shape=target_shape)
    pose_pil = get_pose_pil_by_keypoints(pose_keypoints, target_shape)
    pose_pil = pose_pil.resize(img_pil.size, Image.BILINEAR)

    # Get parse and mask
    parse_pil = get_parse_res(parse_predictor, img_pil=img_pil, target_shape=target_shape)
    mask_pil = get_parse_zone_by_clotype(clo_type=clo_type, parse_pil=parse_pil, target_shape=target_shape)
    mask_pil = mask_pil.resize(img_pil.size, Image.NEAREST)

    # Step 1: masked_img = img.paste(mask, white color)
    masked_img = get_masked_img(img_pil, mask_pil)

    # Step 2: condiction_pose = masked_img.paste(mask, pose_pil)
    cond_pose = get_condition_pose(masked_img, pose_pil, mask_pil)

    # Step 3: condiction_clo = masked_img.paste(mask_rect, clo_resized)
    cond_clo = get_condition_clo(masked_img, mask_pil, clo_pil)

    # Step 4: concat(img, cond_clo, cond_pose)
    final_img = concat_images([img_pil, cond_clo, cond_pose], direction='horizontal')

    return {
        "original": img_pil,
        "masked_img": masked_img,
        "cond_clo": cond_clo,
        "cond_pose": cond_pose,
        "final_img": final_img
    }

if __name__=='__main__':
    # img_dir = '/mnt/nas/shengjie/datasets/DressCode_1024/upper'
    # img_path = os.path.join(img_dir, 'image/000000_0.jpg')
    # clo_path = os.path.join(img_dir, 'cloth/000000_1.jpg')
    # clo_type = 'upper'
    # img_dir = '/mnt/nas/shengjie/datasets/DressCode_1024/lower'
    # img_path = os.path.join(img_dir, 'image/013563_0.jpg')
    # clo_path = os.path.join(img_dir, 'cloth/013563_1.jpg')
    # clo_type = 'lower'
    img_dir = '/mnt/nas/shengjie/datasets/DressCode_1024/dresses'
    img_path = os.path.join(img_dir, 'image/020714_0.jpg')
    clo_path = os.path.join(img_dir, 'cloth/020714_1.jpg')
    clo_type = 'full'
    tryon_pipeline(img_path=img_path,clo_path=clo_path,clo_type=clo_type)
